import torch
from torch.autograd import Function
from torch.autograd.function import once_differentiable
from pcdet.ops.build_dps_cost_volume import build_dps_cost_volume_cuda
# import argparse.ArgumentParser.parse_args as parse_args
from torch.onnx.symbolic_helper import parse_args

class _BuildDpsCostVolume(Function):
    @staticmethod
    def forward(ctx, left, right, shift, psv_channels, downsample, sep=32, interval=1):
        # left（b, 96, 480, 928)
        ctx.save_for_backward(shift, psv_channels)
        ctx.downsample = downsample
        ctx.channels = left.shape[1]
        ctx.sep = sep
        ctx.interval = interval
        assert torch.all(shift >= 0.)
        output = build_dps_cost_volume_cuda.build_dps_cost_volume_forward(
            left, right, shift, psv_channels, downsample, sep, interval)
        return output # [1, 64, 72, 120, 232] shift.size=72 psv_channels.size=65

    @staticmethod
    @once_differentiable
    def backward(ctx, grad_output):
        shift, psv_channels = ctx.saved_tensors
        grad_left, grad_right = build_dps_cost_volume_cuda.build_dps_cost_volume_backward(
            grad_output, shift, psv_channels, ctx.downsample, ctx.channels, ctx.sep, ctx.interval)
        return grad_left, grad_right, None, None, None, None, None

    @staticmethod
    @parse_args('v', 'v', 'v', 'v', 'i', 'i', 'i')
    def symbolic(g, left, right, shift, psv_channels, downsample=4, sep=32, interval=1):
        inputs = [left, right, shift, psv_channels]
        # psv_channels
        kwargs = {
            'downsample_i': downsample,
            'sep_i': sep,
            'interval_i': interval
        }
        # kwargs = [downsample, sep, interval]
        # import pdb; pdb.set_trace()
        return g.op("nvinfer1::DpsCostVolumePlugin", *inputs, **kwargs) 


build_dps_cost_volume = _BuildDpsCostVolume.apply


import torch.onnx.symbolic_registry as sym_registry
def register_onnx_sym_build_dps_cost_volume():
      sym_registry.register_op('nvinfer1::DpsCostVolumePlugin', _BuildDpsCostVolume.symbolic, '', 12)   

register_onnx_sym_build_dps_cost_volume()